import numpy as np
from sklearn.model_selection import KFold
from sklearn.metrics import f1_score
from .model import FastTextModel
import pandas as pd

class Trainer:
    def __init__(self, config, data_processor):
        self.config = config
        self.data_processor = data_processor
        
    def train_and_predict(self, X_train, y_train, X_test):
        """训练模型并进行预测"""
        kf = KFold(n_splits=self.config.N_SPLITS, 
                  random_state=self.config.RANDOM_STATE, 
                  shuffle=True)
        
        test_pred = np.zeros((X_test.shape[0], 1), int)
        
        # 获取完整的训练数据
        train_data = pd.DataFrame({
            'text': X_train,
            'label': y_train,
            'label_ft': '__label__' + y_train.astype(str)
        })
        
        for fold_idx, (train_index, valid_index) in enumerate(kf.split(X_train)):
            print(f'第 {fold_idx+1} 折交叉验证开始...')
            
            # 保存当前折的训练数据，传入train_data和train_index
            self.data_processor.save_temp_train_data(train_data, train_index)
            
            # 训练模型
            model = FastTextModel(self.config)
            model.train(self.config.TEMP_TRAIN_PATH)
            
            # 验证集预测
            val_pred = model.predict(X_train.iloc[valid_index])
            print('Fasttext准确率为：',
                  f1_score(list(y_train.iloc[valid_index]), 
                          val_pred, 
                          average='macro'))
            
            # 测试集预测
            test_pred_ = model.predict(X_test)
            test_pred = np.column_stack((test_pred, test_pred_))
            
        # 取测试集中预测数量最多的数
        final_preds = []
        for i, test_list in enumerate(test_pred):
            final_preds.append(np.argmax(np.bincount(test_list)))
            
        return np.array(final_preds) 